AWS Lambda で SnapStart を有効化した時はアクセスキー取得方法に注意する必要があるようなので検証してみた

いわさです。

AWS Parameters and Secrets Lambda Extension などを使う場合に AWS Lambda で環境変数のAWS_SESSION_TOKENにアクセスする場合があります。

また、Java on Lambda では SnapStart という機能があり、コールドスタート時にスナップショットから復元を行うことで初期化に必要な時間を短縮する機能があります。

先日、環境変数AWS_SESSION_TOKENにアクセスする Lambda 関数の SnapStart を有効化したところ情報が取得出来ませんでした。
実は SnapStart 固有の仕様で、このあたりを考慮する必要があります。

本記事ではセッショントークンの取得方法について検証した結果を紹介します。
ちなみに、AWS_ACCESS_KEY_IDとAWS_SECRET_ACCESS_KEYの場合でも同じです。

SnapStart では環境変数のアクセスキーが設定されない

次のようにAWS_SESSION_TOKEN環境変数へアクセスする Lambda 関数を作成しています。
一応のちほどの確認のためにハンドラー内外で取得してみました。

package helloworld;

import com.amazonaws.services.lambda.runtime.Context;
import com.amazonaws.services.lambda.runtime.RequestHandler;
import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent;
import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyResponseEvent;
import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;

/**
 * Handler for requests to Lambda function.
 */
public class App implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {

    private String hogeXXX = System.getenv("AWS_SESSION_TOKEN");

    public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEvent input, final Context context) {
        APIGatewayProxyResponseEvent response = new APIGatewayProxyResponseEvent();

        String hoge1 = System.getenv("AWS_SESSION_TOKEN");
        String output = String.format("{ \"AWS_SESSION_TOKEN\": \"%s\", \"AWS_SESSION_TOKEN(handler)\": \"%s\"", hogeXXX, hoge1);

        return response
                .withStatusCode(200)
                .withBody(output);
    }
}

SnapStart が有効化されていない場合

SnapStart が有効化されていない Lambda 関数を実行した場合は次のように期待どおりトークンが取得出来ます。

% curl "https://ck98znyec6.execute-api.ap-northeast-1.amazonaws.com/Prod/hello/" | jq
{
  "AWS_SESSION_TOKEN": "hogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehoge",
  "AWS_SESSION_TOKEN(handler)": "hogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehoge"
}

SnapStart を有効化してみると...

今回は SAM で Lambda 関数をデプロイしているので以下のハイライト部分を追加することで SnapStart を有効化した API を作成することが出来ます。
Lambda のコード自体は変更していません。

:

Resources:
  HelloWorldFunction:
    Type: AWS::Serverless::Function # More info about Function Resource: https://github.com/awslabs/serverless-application-model/blob/master/versions/2016-10-31.md#awsserverlessfunction
    Properties:
      CodeUri: HelloWorldFunction
      Handler: helloworld.App::handleRequest
      Runtime: java17
      Architectures:
        - x86_64
      MemorySize: 512
      Environment: # More info about Env Vars: https://github.com/awslabs/serverless-application-model/blob/master/versions/2016-10-31.md#environment-object
        Variables:
          PARAM1: VALUE
      Events:
        HelloWorld:
          Type: Api # More info about API Event Source: https://github.com/awslabs/serverless-application-model/blob/master/versions/2016-10-31.md#api
          Properties:
            Path: /hello
            Method: get
      AutoPublishAlias: SnapStart
      SnapStart:
        ApplyOn: PublishedVersions

Outputs:

:

こちらに再度アクセスしてみます。

% curl "https://ck98znyec6.execute-api.ap-northeast-1.amazonaws.com/Prod/hello/" | jq
{
  "AWS_SESSION_TOKEN": "null",
  "AWS_SESSION_TOKEN(handler)": "null"
}

環境変数に値が設定されていないですね。

SnapStart ではコンテナー認証情報を使用する必要がある

この事象については公式ドキュメントに記述があります。

When SnapStart is activated, the Java runtime automatically uses the container credentials (AWS_CONTAINER_CREDENTIALS_FULL_URI and AWS_CONTAINER_AUTHORIZATION_TOKEN) instead of the access key environment variables. This prevents credentials from expiring before the function is restored.

AWS_CONTAINER_CREDENTIALS_FULL_URIとAWS_CONTAINER_AUTHORIZATION_TOKENを使う必要があると記述されています。
先程の SnapStart を有効化した関数を変更し、これらの値も確認してみましょう。

% curl "https://ck98znyec6.execute-api.ap-northeast-1.amazonaws.com/Prod/hello/" | jq
{
  "AWS_SESSION_TOKEN": "null",
  "AWS_SESSION_TOKEN(handler)": "null",
  "AWS_CONTAINER_CREDENTIALS_FULL_URI": "http://127.0.0.1:9001/2021-04-23/credentials",
  "AWS_CONTAINER_AUTHORIZATION_TOKEN": "bcae7ab7-517c-4604-b2fb-b135868d0c1c"
}

何やら情報が設定されていますね。
この環境変数はコンテナ認証情報プロバイダーで使用されるものです。

AWS SDK for Java のデフォルトクレデンシャルプロバイダーの場合は、環境変数の設定状況から上記の設定値を使って一時アクセス情報を取得するようになっています。

今回のように直接環境変数へアクセスしているような場合と、明示的にクレデンシャルプロバイダーを指定している場合は注意が必要ですね。
AWS_CONTAINER_CREDENTIALS_FULL_URIで取得出来る URL にAWS_CONTAINER_AUTHORIZATION_TOKEN認証情報を付与してリクエストを送信してみます。

:

public class App implements RequestHandler<APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent> {

    public APIGatewayProxyResponseEvent handleRequest(final APIGatewayProxyRequestEvent input, final Context context) {
        APIGatewayProxyResponseEvent response = new APIGatewayProxyResponseEvent();

        String output = getHogeContainerCredential(System.getenv("AWS_CONTAINER_CREDENTIALS_FULL_URI"), System.getenv("AWS_CONTAINER_AUTHORIZATION_TOKEN"));

        return response
                .withStatusCode(200)
                .withBody(output);
    }

    private String getHogeContainerCredential(String uri, String token) {
        try {
            var client = HttpClient.newHttpClient();
            var request = HttpRequest.newBuilder(new URI(uri))
                            .header("Authorization", token)
                            .build();
            final var containerResponse = client.send(request, HttpResponse.BodyHandlers.ofString());

            return containerResponse.body();
        } catch (Exception e) {
            return "hoge";
        }
    }
}

実行してみると次のようにアクセスキー、シークレット、セッショントークンを取得することが出来ました。

% curl "https://ck98znyec6.execute-api.ap-northeast-1.amazonaws.com/Prod/hello/" | jq
{
  "AccessKeyId": "AKIAIOSFODNN7EXAMPLE",
  "SecretAccessKey": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
  "Token": "hogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehogehoge",
  "Expiration": "2023-07-12T08:40:03.249Z"
}

さいごに

本日は AWS Lambda で SnapStart を有効化した時はアクセスキー取得先に注意する必要があるので検証してみました。

このコード変更が必要な旨は Java on AWS Lambda のワークショップにも実は記述があります。

However with SnapStart enabled the Java runtime switches to using container credentials (AWS_CONTAINER_CREDENTIALS_FULL_URI and AWS_CONTAINER_AUTHORIZATION_TOKEN) instead of the access key environment variables.

今回は HTTP クライアントからリクエストしましたが、クレデンシャルプロバイダーにまかせてそこから必要な情報を取得するとか、もう少しうまい方法あるかもしれないなと思いました。

なお今回の方法を使って、AWS Parameters and Secrets Lambda Extension の利用を検証したわけではないので、そのあたりは試してもらえたらと思います。